/**************************************************************************
**
** Copyright (c) 2011 visualfc. All rights reserved.
**
** This library is free software; you can redistribute it and/or
** modify it under the terms of the GNU Lesser General Public
** License as published by the Free Software Foundation; either
** version 2.1 of the License, or (at your option) any later version.
**
** This library is distributed in the hope that it will be useful,
** but WITHOUT ANY WARRANTY; without even the implied warranty of
** MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
** Lesser General Public License for more details.
**
**************************************************************************/
// Module: zmq
// Creator: visualfc <visualfc@gmail.com>
// date: 2011-6-15
// $Id: zmq.go,v 1.0 2011-8-17 visualfc Exp $

package zmq

//#cgo LDFLAGS: -lzmq
//#include "gozmq.h"
import "C"

import (
	"unsafe"
	"os"
	"fmt"
	"sync"
	"time"
	"io"
	"reflect"
)

var newTimerMutex sync.Mutex

func Sleep(ns int64) {
	newTimerMutex.Lock()
	t := time.NewTimer(ns)
	newTimerMutex.Unlock()
	<-t.C
}

// Socket type
type SocketType int
// Send/Recv flag
type TransFlag int
// Socket option
type SocketOption int
// Pollin event flag
type EventFlag int
// device type
type DeviceType int

const (
	// SocketType
	PAIR   SocketType = C.ZMQ_PAIR
	PUB    SocketType = C.ZMQ_PUB
	SUB    SocketType = C.ZMQ_SUB
	REQ    SocketType = C.ZMQ_REQ
	REP    SocketType = C.ZMQ_REP
	DEALER SocketType = C.ZMQ_DEALER
	ROUTER SocketType = C.ZMQ_ROUTER
	PULL   SocketType = C.ZMQ_PULL
	PUSH   SocketType = C.ZMQ_PUSH
	XPUB   SocketType = C.ZMQ_XPUB
	XSUB   SocketType = C.ZMQ_XSUB
	// SocketOption
	HWM               SocketOption = C.ZMQ_HWM
	SWAP              SocketOption = C.ZMQ_SWAP
	AFFINITY          SocketOption = C.ZMQ_AFFINITY
	IDENTITY          SocketOption = C.ZMQ_IDENTITY
	SUBSCRIBE         SocketOption = C.ZMQ_SUBSCRIBE
	UNSUBSCRIBE       SocketOption = C.ZMQ_UNSUBSCRIBE
	RATE              SocketOption = C.ZMQ_RATE
	RECOVERY_IVL      SocketOption = C.ZMQ_RECOVERY_IVL
	MCAST_LOOP        SocketOption = C.ZMQ_MCAST_LOOP
	SNDBUF            SocketOption = C.ZMQ_SNDBUF
	RCVBUF            SocketOption = C.ZMQ_RCVBUF
	RCVMORE           SocketOption = C.ZMQ_RCVMORE
	FD                SocketOption = C.ZMQ_FD
	EVENTS            SocketOption = C.ZMQ_EVENTS
	TYPE              SocketOption = C.ZMQ_TYPE
	LINGER            SocketOption = C.ZMQ_LINGER
	RECONNECT_IVL     SocketOption = C.ZMQ_RECONNECT_IVL
	BACKLOG           SocketOption = C.ZMQ_BACKLOG
	RECOVERY_IVL_MSEC SocketOption = C.ZMQ_RECOVERY_IVL_MSEC /*  opt. recovery time, reconcile in 3.x   */
	RECONNECT_IVL_MAX SocketOption = C.ZMQ_RECONNECT_IVL_MAX
	// TransFlag
	DONTWAIT TransFlag = C.ZMQ_DONTWAIT
	NOBLOCK  TransFlag = DONTWAIT
	SNDMORE  TransFlag = C.ZMQ_SNDMORE
	SNDLABEL TransFlag = C.ZMQ_SNDLABEL
	// PollinType
	POLLIN  EventFlag = C.ZMQ_POLLIN
	POLLOUT EventFlag = C.ZMQ_POLLOUT
	POLLERR EventFlag = C.ZMQ_POLLERR
	// DeviceType
	STREAMER  DeviceType = C.ZMQ_STREAMER
	FORWARDER DeviceType = C.ZMQ_FORWARDER
	QUEUE     DeviceType = C.ZMQ_QUEUE
)

// zmq error
type Error struct {
	Errno int
	Text  string
}

func (e *Error) String() string {
	return fmt.Sprintf("%s, %d, %s", e.Text, e.Errno, C.GoString(C.zmq_strerror(C.int(e.Errno))))
}

// send DONTWAIT mode : if the message cannot be queued on the socket, return ErrAgain
// recv DONTWAIT mode : if there are no bytes available return nil,ErrAgain
var (
	ErrAgain = &Error{Errno: int(C.EAGAIN), Text: "DONTWAIT"}
)

// zmq os.Error
func zmqError(text string) os.Error {
	return &Error{Errno: int(C.zmq_errno()), Text: text}
}

// zmq Version
type Ver struct {
	Major, Minor, Patch int
}

func (v *Ver) String() string {
	return fmt.Sprintf("%d.%d.%d", v.Major, v.Minor, v.Patch)
}

// ØMQ version
func Version() *Ver {
	var major, minor, patch C.int
	C.zmq_version(&major, &minor, &patch)
	return &Ver{int(major), int(minor), int(patch)}
}

// zmq context
type Context struct {
	ctx unsafe.Pointer
}

// zmq message
type Msg struct {
	msg *C.zmq_msg_t
}

// zmq socket
type Socket struct {
	sock unsafe.Pointer
}

// initialise ØMQ context 
func Init(io_threads int) (*Context, os.Error) {
	ctx := C.zmq_init(C.int(io_threads))
	if ctx == nil {
		return nil, zmqError("zmq_init")
	}
	return &Context{ctx}, nil
}

// check ØMQ context is valid
func (c *Context) IsValid() bool {
	return c != nil && c.ctx != nil
}

// terminate ØMQ context 
func (c *Context) Close() os.Error {
	if c == nil || c.ctx == nil {
		return os.NewError("nil zmq context")
	}
	if C.zmq_term(c.ctx) != 0 {
		return zmqError("zmq_term")
	}
	c.ctx = nil
	return nil
}

// initialise empty ØMQ message 
func InitMsg() (*Msg, os.Error) {
	m := new(Msg)
	m.msg = new(C.zmq_msg_t)
	if C.zmq_msg_init(m.msg) != 0 {
		return nil, zmqError("zmq_msg_init")
	}
	return m, nil
}

// initialise ØMQ message from a supplied buffer 
func InitMsgWithData(data []byte) (*Msg, os.Error) {
	m := new(Msg)
	m.msg = new(C.zmq_msg_t)
	size := len(data)
	if size == 0 {
		if C.zmq_msg_init(m.msg) != 0 {
			return nil, zmqError("zmq_msg_init")
		}
	} else {
		if C.zmq_msg_init_data_helper(m.msg, unsafe.Pointer(&data[0]), C.size_t(len(data))) != 0 {
			return nil, zmqError("zmq_msg_init_data")
		}
	}
	return m, nil
}

// initialise ØMQ message of a specified size 
func InitMsgWithSize(size int) (*Msg, os.Error) {
	m := new(Msg)
	m.msg = new(C.zmq_msg_t)
	if C.zmq_msg_init_size(m.msg, C.size_t(size)) != 0 {
		return nil, zmqError("zmq_msg_init_size")
	}
	return m, nil
}

// writes the contents of p into the message. 
// It returns the number of bytes written. 
// If nn < m.Size(), it also returns an error explaining why the write is short. 
func (m *Msg) Write(p []byte) (nn int, err os.Error) {
	size := len(p)
	if size == 0 {
		return 0, nil
	}
	sz := int(C.zmq_msg_size(m.msg))
	if size > sz {
		size = sz
	}
	C.memcpy(unsafe.Pointer(C.zmq_msg_data(m.msg)), unsafe.Pointer(&p[0]), C.size_t(size))
	return size, nil
}

// check ØMQ msg is valid
func (m *Msg) IsValid() bool {
	return m != nil && m.msg != nil
}

// close ØMQ msg
func (m *Msg) Close() os.Error {
	if m == nil || m.msg == nil {
		return os.NewError("nil zmq msg")
	}
	if C.zmq_msg_close(m.msg) != 0 {
		return zmqError("zmq_msg_close")
	}
	m.msg = nil
	return nil
}

func (m *Msg) Size() int {
	return int(C.zmq_msg_size(m.msg))
}

func (m *Msg) Data() []byte {
	var data []byte
	size := C.zmq_msg_size(m.msg)
	if size > 0 {
		data = (*[1 << 30]byte)(unsafe.Pointer(C.zmq_msg_data(m.msg)))[0:size]
	}
	return data
}

func (m *Msg) Copy(src *Msg) os.Error {
	if C.zmq_msg_copy(m.msg, src.msg) != 0 {
		return zmqError("zmq_msg_copy")
	}
	return nil
}

func (m *Msg) Move(src *Msg) os.Error {
	if C.zmq_msg_move(m.msg, src.msg) != 0 {
		return zmqError("zmq_msg_move")
	}
	return nil
}

// create ØMQ socket 
func (c *Context) Socket(socket_type SocketType) (*Socket, os.Error) {
	sock := C.zmq_socket(c.ctx, C.int(socket_type))
	if sock == nil {
		return nil, zmqError("zmq_socket")
	}
	return &Socket{sock}, nil
}

// check ØMQ socket is valid
func (s *Socket) IsValid() bool {
	return s != nil && s.sock != nil
}

// close ØMQ socket 
func (s *Socket) Close() os.Error {
	if s == nil || s.sock == nil {
		return os.NewError("nil zmq socket")
	}
	if C.zmq_close(s.sock) != 0 {
		return zmqError("zmq_close")
	}
	s.sock = nil
	return nil
}

func (s *Socket) getOptionValue(option SocketOption, val unsafe.Pointer, size C.size_t) (int, os.Error) {
	if C.zmq_getsockopt(s.sock, C.int(option), val, &size) != 0 {
		return 0, zmqError(fmt.Sprintf("zmq_getsockopt %d", option))
	}
	return int(size), nil
}

func (s *Socket) GetOptionVal(option SocketOption, val interface{}) os.Error {
	switch v := (val).(type) {
	case *int:
		_, err := s.getOptionValue(option, unsafe.Pointer(val.(*int)), C.size_t(unsafe.Sizeof(val)))
		return err
	case *int32:
		_, err := s.getOptionValue(option, unsafe.Pointer(val.(*int32)), C.size_t(unsafe.Sizeof(val)))
		return err
	case *uint32:
		_, err := s.getOptionValue(option, unsafe.Pointer(val.(*uint32)), C.size_t(unsafe.Sizeof(val)))
		return err
	case *int64:
		_, err := s.getOptionValue(option, unsafe.Pointer(val.(*int64)), C.size_t(unsafe.Sizeof(val)))
		return err
	case *uint64:
		_, err := s.getOptionValue(option, unsafe.Pointer(val.(*uint64)), C.size_t(unsafe.Sizeof(val)))
		return err
	case *[]byte:
		s, err := s.GetOptionByte(option)
		if err == nil {
			*v = s
		}
		return err
	case *string:
		s, err := s.GetOptionString(option)
		if err == nil {
			*v = s
		}
		return err
	default:
		return os.NewError("unsupported type : " + reflect.TypeOf(v).String())
	}
	return os.NewError("unkown error")
}

func (s *Socket) GetOptionInt(option SocketOption) (int, os.Error) {
	var val int
	_, err := s.getOptionValue(option, unsafe.Pointer(&val), C.size_t(unsafe.Sizeof(val)))
	return val, err
}

func (s *Socket) GetOptionInt32(option SocketOption) (int32, os.Error) {
	var val int32
	_, err := s.getOptionValue(option, unsafe.Pointer(&val), C.size_t(unsafe.Sizeof(val)))
	return val, err
}

func (s *Socket) GetOptionUInt32(option SocketOption) (uint32, os.Error) {
	var val uint32
	_, err := s.getOptionValue(option, unsafe.Pointer(&val), C.size_t(unsafe.Sizeof(val)))
	return val, err
}

func (s *Socket) GetOptionInt64(option SocketOption) (int64, os.Error) {
	var val int64
	_, err := s.getOptionValue(option, unsafe.Pointer(&val), C.size_t(unsafe.Sizeof(val)))
	return val, err
}

func (s *Socket) GetOptionUInt64(option SocketOption) (uint64, os.Error) {
	var val uint64
	_, err := s.getOptionValue(option, unsafe.Pointer(&val), C.size_t(unsafe.Sizeof(val)))
	return val, err
}

func (s *Socket) GetOptionByte(option SocketOption) ([]byte, os.Error) {
	val := make([]byte, 1024)
	sz, err := s.getOptionValue(option, unsafe.Pointer(&val[0]), C.size_t(len(val)))
	if sz == 0 {
		return nil, err
	}
	return val[:sz], err
}

func (s *Socket) GetOptionString(option SocketOption) (string, os.Error) {
	val, err := s.GetOptionByte(option)
	return string(val), err
}

func (s *Socket) setOptionValue(option SocketOption, val unsafe.Pointer, size C.size_t) os.Error {
	if C.zmq_setsockopt(s.sock, C.int(option), val, size) != 0 {
		return zmqError(fmt.Sprintf("zmq_setsockopt %d", option))
	}
	return nil
}

func (s *Socket) SetOptionVal(option SocketOption, val interface{}) os.Error {
	switch v := val.(type) {
	case int:
		return s.SetOptionInt(option, val.(int))
	case int32:
		return s.SetOptionInt32(option, val.(int32))
	case uint32:
		return s.SetOptionUInt32(option, val.(uint32))
	case int64:
		return s.SetOptionInt64(option, val.(int64))
	case uint64:
		return s.SetOptionUInt64(option, val.(uint64))
	case []byte:
		return s.SetOptionByte(option, val.([]byte))
	case string:
		return s.SetOptionString(option, val.(string))
	default:
		return os.NewError("unsupported type : " + reflect.TypeOf(v).String())
	}
	return os.NewError("unkown error")
}

func (s *Socket) SetOptionInt(option SocketOption, val int) os.Error {
	return s.setOptionValue(option, unsafe.Pointer(&val), C.size_t(unsafe.Sizeof(val)))
}

func (s *Socket) SetOptionInt32(option SocketOption, val int32) os.Error {
	return s.setOptionValue(option, unsafe.Pointer(&val), C.size_t(unsafe.Sizeof(val)))
}

func (s *Socket) SetOptionUInt32(option SocketOption, val uint32) os.Error {
	return s.setOptionValue(option, unsafe.Pointer(&val), C.size_t(unsafe.Sizeof(val)))
}

func (s *Socket) SetOptionInt64(option SocketOption, val int64) os.Error {
	return s.setOptionValue(option, unsafe.Pointer(&val), C.size_t(unsafe.Sizeof(val)))
}

func (s *Socket) SetOptionUInt64(option SocketOption, val uint64) os.Error {
	return s.setOptionValue(option, unsafe.Pointer(&val), C.size_t(unsafe.Sizeof(val)))
}

func (s *Socket) SetOptionByte(option SocketOption, val []byte) os.Error {
	v := C.CString(string(val))
	defer C.free(unsafe.Pointer(v))
	return s.setOptionValue(option, unsafe.Pointer(v), C.size_t(len(val)))
}

func (s *Socket) SetOptionString(option SocketOption, val string) os.Error {
	return s.SetOptionByte(option, []byte(val))
}

func (s *Socket) GetRecvMore() (int64, os.Error) {
	return s.GetOptionInt64(RCVMORE)
}

// accept connections on a socket 
func (s *Socket) Bind(addr string) os.Error {
	a := C.CString(addr)
	defer C.free(unsafe.Pointer(a))
	if C.zmq_bind(s.sock, a) != 0 {
		return zmqError("zmq_bind " + addr)
	}
	return nil
}

// connect a socket 
func (s *Socket) Connect(addr string) os.Error {
	a := C.CString(addr)
	defer C.free(unsafe.Pointer(a))
	if C.zmq_connect(s.sock, a) != 0 {
		return zmqError("zmq_connect " + addr)
	}
	return nil
}

// send a bytes on a socket 
// flags - DONTWAIT mode : if the message cannot be queued on the socket, return ErrAgain
func (s *Socket) Send(data []byte, flags TransFlag) os.Error {
	var m C.zmq_msg_t
	size := len(data)
	if size == 0 {
		if C.zmq_msg_init(&m) != 0 {
			return zmqError("zmq_msg_init")
		}
	} else {
		if C.zmq_msg_init_data_helper(&m, unsafe.Pointer(&data[0]), C.size_t(size)) != 0 {
			return zmqError("zmq_msg_init_data")
		}
	}
	defer C.zmq_msg_close(&m)
	if C.zmq_send_msg(s.sock, &m, C.int(flags)) != 0 {
		if flags&DONTWAIT == DONTWAIT && C.zmq_errno() == C.EAGAIN {
			return ErrAgain
		}
		return zmqError("zmq_send")
	}
	return nil
}

// send a string on a socket
func (s *Socket) SendString(str string, flags TransFlag) os.Error {
	return s.Send([]byte(str), flags)
}

// send all bytes list on a socket, auto append SNDMORE 
func (s *Socket) SendAll(data ...[]byte) os.Error {
	size := len(data)
	if size == 0 {
		return s.Send(nil,0)
	}
	for i := 0; i < size-1; i++ {
		err := s.Send(data[i],SNDMORE)
		if err != nil {
			return err
		}
	}
	return s.Send(data[size-1],0)
}

// send all string list on a socket, auto append SNDMORE 
func (s *Socket) SendAllString(str ...string) os.Error {
	size := len(str)
	if size == 0 {
		return s.Send(nil,0)
	}
	for i := 0; i < size-1; i++ {
		err := s.Send([]byte(str[i]),SNDMORE)
		if err != nil {
			return err
		}
	}
	return s.Send([]byte(str[size-1]),0)
}

// send a message on a socket 
// flags - DONTWAIT mode : if the message cannot be queued on the socket, return ErrAgain
func (s *Socket) SendMsg(m *Msg, flags TransFlag) os.Error {
	if C.zmq_send_msg(s.sock, m.msg, C.int(flags)) != 0 {
		if flags&DONTWAIT == DONTWAIT && C.zmq_errno() == C.EAGAIN {
			return ErrAgain
		}
		return zmqError("zmq_send_msg")
	}
	return nil
}

// receive a bytes from a socket 
// return recv message bytes and os.Error
// flags - ONTWAIT mode : if there are no bytes available return nil,ErrAgain
func (s *Socket) Recv(flags TransFlag) ([]byte, os.Error) {
	var m C.zmq_msg_t
	if C.zmq_msg_init(&m) != 0 {
		return nil, zmqError("zmq_msg_init")
	}
	defer C.zmq_msg_close(&m)
	if C.zmq_recv_msg(s.sock, &m, C.int(flags)) != 0 {
		if flags&DONTWAIT == DONTWAIT && C.zmq_errno() == C.EAGAIN {
			return nil, ErrAgain
		}
		return nil, zmqError("zmq_recv")
	}
	size := C.zmq_msg_size(&m)
	if size == 0 {
		return nil, nil
	}
	data := make([]byte, int(size))
	C.memcpy(unsafe.Pointer(&data[0]), C.zmq_msg_data(&m), size)
	return data, nil
}

// receive a string from a socket
func (s *Socket) RecvString(flags TransFlag) (string, os.Error) {
	data, err := s.Recv(flags)
	return string(data), err
}

// receive a message from a socket 
// return recv message bytes and os.Error
// flags - DONTWAIT mode : if there are no messages available return nil,ErrAgain
func (s *Socket) RecvMsg(m *Msg, flags TransFlag) ([]byte, os.Error) {
	if C.zmq_recv_msg(s.sock, m.msg, C.int(flags)) != 0 {
		if flags&DONTWAIT == DONTWAIT && C.zmq_errno() == C.EAGAIN {
			return nil, ErrAgain
		}
		return nil, zmqError("zmq_recv_msg")
	}
	return m.Data(), nil
}

func (s *Socket) Dump() os.Error {
	return s.Fdump(os.Stdout)
}

func (s *Socket) Fdump(w io.Writer) os.Error {
	fmt.Fprintln(w, "----------------------------------------")
	msg, err := InitMsg()
	if err != nil {
		return err
	}
	defer msg.Close()

	for {
		data, err := s.RecvMsg(msg, 0)
		if err != nil {
			return err
		}
		size := msg.Size()
		is_text := true

		for _, v := range data {
			if v < 32 || v > 127 {
				is_text = false
				break
			}
		}
		fmt.Fprintf(w, "[%03d] ", size)
		if is_text {
			fmt.Fprintf(w, "%s\n", data)
		} else {
			fmt.Fprintf(w, "%X\n", data)
		}
		more, err := s.GetOptionInt64(RCVMORE)
		if err != nil {
			return err
		}
		if more == 0 {
			break
		}
	}
	return nil
}

func Device(device DeviceType, in, out *Socket) os.Error {
	if in == nil || in.sock == nil || out == nil || out.sock == nil {
		return os.NewError("nil zmq socket")
	}
	if C.zmq_device(C.int(device), in.sock, out.sock) != 0 {
		return zmqError("zmq_device")
	}
	return nil
}

// ØMQ PollItem
type PollItem C.zmq_pollitem_t

func (item *PollItem) SetSocket(s *Socket) {
	item.socket = s.sock
}

func (item *PollItem) SetEvents(flags EventFlag) {
	item.events = C.short(flags)
}

func (item *PollItem) GetEvents() EventFlag {
	return EventFlag(item.events)
}

func (item *PollItem) GetRevents() EventFlag {
	return EventFlag(item.revents)
}

// Upon successful completion, the zmq.Poll() function shall return the number 
// of PollItem structures with events signaled in revents or 0 if no events 
// have been signaled. Upon failure, zmq.Poll() shall return -1 and os.Error
func Poll(items []PollItem, timeout int64) (n int, err os.Error) {
	size := len(items)
	if size == 0 {
		err = os.NewError("nil zmq poll items")
		return
	}
	n = int(C.zmq_poll((*C.zmq_pollitem_t)(&items[0]), C.int(size), C.long(timeout)))
	if n == -1 {
		err = zmqError("zmq_poll")
	}
	return
}
